import os
import random
import shutil
import json
import zipfile

'''
实现将图片文件和对应的标注文件随机分配到多个子文件夹，并且平均分配：
从文件夹A中创建指定名称的多个子文件夹。
从文件夹B中读取所有图片文件和对应的标注文件。
将图片文件和对应的标注文件按照平均分配的方式随机分配到子文件夹中。
'''

# 父文件夹A中的指定名称列表a
subfolder_names = ["杨健凯",
     "汤小婷",
     "谭紫秋",
     "罗炳先",
     "徐冉",
     "朱彦奎",
     "杨剑",
     "吴梓鹏",
     "刘梦杨",
     "罗思思",
     "蒋依柔",
     "高皓天",
     "李浩淼",
     "黄聪利",
     "凌子敬",
     "韦震天",
     "潘永强",
     "韦日翔",
     "刘钦忠"]
# 文件夹A和文件夹B的路径
folder_a_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/分批"
folder_b_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/good all"


# 创建子文件夹的函数
def create_subfolders(parent_folder, subfolder_names):
    subfolders = []
    for folder_name in subfolder_names:
        folder_path = os.path.join(parent_folder, folder_name)
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        subfolders.append(folder_path)
    return subfolders

# 将文件分配到子文件夹的函数
def distribute_files(parent_folder_a, parent_folder_b, subfolders):
    all_files = [f for f in os.listdir(parent_folder_b) if f.endswith(('.jpg', '.png'))]  # 假设是图片格式
    labeled_images = []
    unlabeled_images = []

    # 分离有标注和无标注的图片
    for image in all_files:
        image_path = os.path.join(parent_folder_b, image)
        label_path = os.path.join(parent_folder_b, f"{os.path.splitext(image)[0]}.json")
        if os.path.exists(label_path):  # 如果有标注文件
            labeled_images.append(image)
        else:  # 如果没有标注文件
            unlabeled_images.append(image)

    # 随机分配有标注文件的图片到子文件夹
    random.shuffle(labeled_images)
    for idx, image in enumerate(labeled_images):
        subfolder = subfolders[idx % len(subfolders)]
        image_path = os.path.join(parent_folder_b, image)
        shutil.copy(image_path, os.path.join(subfolder, image))
        # 复制标注文件
        label_file = f"{os.path.splitext(image)[0]}.json"
        shutil.copy(os.path.join(parent_folder_b, label_file), os.path.join(subfolder, label_file))

    # 随机分配没有标注文件的图片到子文件夹
    random.shuffle(unlabeled_images)
    for idx, image in enumerate(unlabeled_images):
        subfolder = subfolders[idx % len(subfolders)]
        image_path = os.path.join(parent_folder_b, image)
        shutil.copy(image_path, os.path.join(subfolder, image))

# 创建子文件夹
subfolders = create_subfolders(folder_a_path, subfolder_names)



# 分配文件
distribute_files(folder_a_path, folder_b_path, subfolders)
print("文件已成功分配到子文件夹！")

def compress_first_subfolders(source_dir, output_dir):
    for root, dirs, files in os.walk(source_dir):
        for dir_name in dirs:
            subfolder_path = os.path.join(root, dir_name)

            # 获取第一个子文件夹
            if subfolder_path == source_dir or subfolder_path.startswith(source_dir + os.sep):
                zip_file_name = os.path.join(output_dir, f"{dir_name}.zip")

                # 创建一个新的 zip 文件并将子文件夹压缩进去
                with zipfile.ZipFile(zip_file_name, 'w') as zipf:
                    for root_sub, dirs_sub, files_sub in os.walk(subfolder_path):
                        for file_sub in files_sub:
                            file_path = os.path.join(root_sub, file_sub)
                            zipf.write(file_path, os.path.relpath(file_path, source_dir))

# 压缩第一个子文件夹
compress_first_subfolders(folder_a_path, folder_a_path)

